#!/usr/bin/env python3
"""
An utility to convert source models into ruptures
"""
import os
import logging
import operator
import numpy
from openquake.baselib import sap, parallel, general, hdf5
from openquake.calculators import base
from openquake.calculators.base import expose_outputs

rupfields = dict(
    mag=numpy.float32,
    occurrence_rate=numpy.float32,
    hypo_lon=numpy.float32,
    hypo_lat=numpy.float32,
    hypo_dep=numpy.float32,
    grp_id=numpy.uint16,
    rup_id=numpy.uint32,
    src_id=numpy.uint32,
    # add more if you like
)


# NB: this is called after a preclassical calculation, so area sources and
# multipoint sources are already split
def build_rupdict(srcs):
    """
    :param srcs: a list of sources of the same source group
    """
    dic = {f: [] for f in rupfields}
    for src in srcs:
        irups = src.iter_ruptures(step=5)  # fast
        for rupno, rup in enumerate(irups):
            dic['mag'].append(rup.mag)
            dic['occurrence_rate'].append(rup.occurrence_rate)
            dic['hypo_lon'].append(rup.hypocenter.x)
            dic['hypo_lat'].append(rup.hypocenter.y)
            dic['hypo_dep'].append(rup.hypocenter.z)
            dic['grp_id'].append(src.grp_id)
            dic['src_id'].append(src.id)
            dic['rup_id'].append(src.offset + rupno)
    return {k: rupfields[k](v) for k, v in dic.items()}


def csm2rup(csm, dstore):
    """
    :param csm: CompositeSourceModel instance
    :param dstore: DataStore where to save the ruptures
    """
    def src_weight(src):
        # make point sources 100x lighter, due to iter_ruptures_fast
        return .01 if src.code == b'P' else 1

    sources = csm.get_sources()
    dstore.create_df('rup', rupfields.items(), compression='gzip')  # 10x gain!
    for rupdict in parallel.Starmap.apply(
            build_rupdict, (sources,), weight=src_weight,
            key=operator.attrgetter('grp_id'), h5=dstore,
            concurrent_tasks=1000):
        for k, v in rupdict.items():
            hdf5.extend(dstore[f'rup/{k}'], v)
    nrups = len(dstore['rup/mag'])
    size = general.humansize(os.path.getsize(dstore.filename))
    logging.info('Stored {:_d} ruptures [{}] inside {}'.format(
        nrups, size, dstore.filename))


def main(job_ini):
    """
    Convert a composite source model into a DataFrame of ruptures
    """
    calc = base.run_calc(job_ini, calculation_mode='preclassical')
    csm2rup(calc.csm, calc.datastore)
    expose_outputs(calc.datastore)

if __name__ == '__main__':
    sap.run(main)
